'''
从B3DM中提取出GLTF，然后提取出纹理和纹理MASK图
'''
from filecmp import cmp
import json
from pathlib import Path
import unittest

import numpy as np
import sys 
sys.path.append("./") 

# from py3dtiles.tileset.content import B3dm, GlTF, read_binary_tile_content

# from py3dtiles.tileset import B3dm, GlTF, read_binary_tile_content
from py3dtiles import TileContentReader as TileContentReader
import glob
import struct
import trimesh
import cv2
from tqdm import tqdm
import os

datadir = '/mnt/B0AAF35FAAF32110/data/3DTilesVPR/test2/'
file = 'Tile_+000_+000_L22_00013300.b3dm'

def extract_gltf(path):
    # tile_content = read_binary_tile_content(
    #     Path(path)
    # )  # TODO re-export b3dm once feature table is added
    # glt = tile_content.body.gltf

    reader =  TileContentReader()
    tile_content = reader.read_file(Path(path))
    glt = tile_content.body.glTF
    arr = glt.to_array()
    gltffile = open(path + ".glb" , "wb")
    gltffile.write(arr)

def extract_image(path):
    bottle_gltf = trimesh.load(path)
    # bottle_trimesh = bottle_gltf.geometry[list(bottle_gltf.geometry.keys())[0]]

    gltf_mesh = bottle_gltf.geometry['GLTF']
    cvImg = None
    pilImg = gltf_mesh.visual.material.baseColorTexture
    if pilImg is None:
        print("None Image:"+ path)
        return
    cvImg = cv2.cvtColor(np.array(pilImg),cv2.COLOR_RGB2BGR)

    width = cvImg.shape[1]
    height =cvImg.shape[0]

    uvAll = gltf_mesh.visual.uv
    facesAll = gltf_mesh.faces
    facesNormalAll = gltf_mesh.face_normals

    uvImg = np.zeros(cvImg.shape,dtype=cvImg.dtype)

    for i in range(facesAll.shape[0]):
        face = facesAll[i]
        
        
        uv_one = uvAll[face]
        uv_one[:,0] *= width
        uv_one[:,1] = height*(1-uv_one[:,1])
        uv_one = np.array(uv_one).astype('int32')
        
        # normal = facesNormalAll[i]
        # color = normal*255
        
        cv2.fillPoly(uvImg,[uv_one],(255,255,255))
        # cv2.polylines(uvImg,[uv_one],True,(0,255,0))
    
    cv2.imwrite(path + ".jpg",cvImg)
    cv2.imwrite(path + ".MASK.png",uvImg)


# 1 提取GLTF文件
files = glob.glob(datadir + "database/*/*.b3dm",recursive=True)

pbar = tqdm(total=len(files))
for file in files:
    pbar.update()
    if os.path.exists(Path(file + ".glb")):
        continue
    extract_gltf(file)
pbar.close()

# files = glob.glob(datadir + "*/*.glb")
# pbar = tqdm(total=len(files))
# for file in files:
#     extract_image(file)
#     pbar.update()zzdx_B3DM
# pbar.close()

files = glob.glob(datadir + "database/*/*.glb")
pbar = tqdm(total=len(files))
for file in files:
    pbar.update()
    if os.path.exists(Path(file + ".jpg")):
        continue
    extract_image(file)
    
pbar.close()